﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.Utils;
using Algoritmia.DataStructures.PriorityMaps;
using Algoritmia.SemiRings;
using Algoritmia.StateSpaces;

namespace Algoritmia.Schemes.DynamicProgramming
{
    public class ForwardWithPriorityQueueDynamicProgrammingSolver<TState,TDecision,TResult,TWeight>
    {
        private readonly Func<IEnumerable<KeyValuePair<TState, TResult>>, IPriorityMap<TState, TResult>> _createPriorityMap;

        public ForwardWithPriorityQueueDynamicProgrammingSolver(Func<IEnumerable<KeyValuePair<TState, TResult>>, IPriorityMap<TState, TResult>> createPriorityMap = null)
        {
            var comparison = AlgoritmiaFactory.GetDefaultComparison<TResult>();
            _createPriorityMap = createPriorityMap ?? (it => new HeapMap<TState, TResult>(it,firstHasHigherPriority: (kv1,kv2) => comparison(kv1.Value,kv2.Value) < 0 ));
        }

        public TResult Solve(IForwardStateSpace<TState,TDecision> space, ISemiRing<TResult,TWeight> semiRing,
            Func<TState, TDecision, TWeight> wfn)
        {
            var mem = _createPriorityMap(space.InitialStates.Select(i => new KeyValuePair<TState,TResult>(i,semiRing.One)));
            var bestFinalScore = semiRing.Zero;
            while ( mem.Count > 0)
            {
                var optimal = mem.ExtractOptimal();
                var s = optimal.Key;
                var score = optimal.Value;
                foreach(var d in space.Decisions(s))
                {
                    var s2 = space.Decide(s, d);
                    if ( !mem.Contains(s2))
                    {
                        mem[s2] = semiRing.Zero;
                    }
                    mem[s2] = semiRing.Plus(mem[s2], semiRing.Times(score, wfn(s, d)));
                }
                if ( space.IsFinal(s))
                {
                    bestFinalScore = semiRing.Plus(bestFinalScore, score);
                }
            }
            return bestFinalScore;
        }
    }

    public class ForwardWithPriorityQueueDynamicProgrammingSolver<TState, TDecision, TResult> :
        ForwardWithPriorityQueueDynamicProgrammingSolver<TState, TDecision, TResult, TResult>
    {
        public ForwardWithPriorityQueueDynamicProgrammingSolver(Func<IEnumerable<KeyValuePair<TState, TResult>>, IPriorityMap<TState, TResult>> createPriorityMap = null)
            : base(createPriorityMap)
        {
        }
    }
}
